package AdvancedEncryptionStandard;

import org.apache.commons.codec.binary.Hex;

public class AES {
    private static final byte[][] SBOXs = new byte[16][16];
    private static final String[][] SBOXsStr = new String[][]{
            {"63", "7C", "77", "7B", "F2", "6B", "6F", "C5", "30", "01", "67", "2B", "FE", "D7", "AB", "76"},
            {"CA", "82", "C9", "7D", "FA", "59", "47", "F0", "AD", "D4", "A2", "AF", "9C", "A4", "72", "C0"},
            {"B7", "FD", "93", "26", "36", "3F", "F7", "CC", "34", "A5", "E5", "F1", "71", "D8", "31", "15"},
            {"04", "C7", "23", "C3", "18", "96", "05", "9A", "07", "12", "80", "E2", "EB", "27", "B2", "75"},
            {"09", "83", "2C", "1A", "1B", "6E", "5A", "A0", "52", "3B", "D6", "B3", "29", "E3", "2F", "84"},
            {"53", "D1", "00", "ED", "20", "FC", "B1", "5B", "6A", "CB", "BE", "39", "4A", "4C", "58", "CF"},
            {"D0", "EF", "AA", "FB", "43", "4D", "33", "85", "45", "F9", "02", "7F", "50", "3C", "9F", "A8"},
            {"51", "A3", "40", "8F", "92", "9D", "38", "F5", "BC", "B6", "DA", "21", "10", "FF", "F3", "D2"},
            {"CD", "0C", "13", "EC", "5F", "97", "44", "17", "C4", "A7", "7E", "3D", "64", "5D", "19", "73"},
            {"60", "81", "4F", "DC", "22", "2A", "90", "88", "46", "EE", "B8", "14", "DE", "5E", "0B", "D8"},
            {"E0", "32", "3A", "0A", "49", "06", "24", "5C", "C2", "D3", "AC", "62", "91", "95", "E4", "79"},
            {"E7", "C8", "37", "6D", "8D", "D5", "4E", "A9", "6C", "56", "F4", "EA", "65", "7A", "AE", "08"},
            {"BA", "78", "25", "2E", "1C", "A6", "B4", "C6", "E8", "DD", "74", "1F", "4B", "BD", "8B", "8A"},
            {"70", "3E", "B5", "66", "48", "03", "F6", "0E", "61", "35", "57", "B9", "86", "C1", "1D", "9E"},
            {"E1", "F8", "98", "11", "69", "D9", "8E", "94", "9B", "1E", "87", "E9", "CE", "55", "28", "DF"},
            {"8C", "A1", "89", "0D", "BF", "E6", "42", "68", "41", "99", "2D", "0F", "B0", "54", "BB", "16"}
    };

    private static final byte[][] InverseSBOXs = new byte[16][16];
    private static final String[][] InverseSBOXsStr = new String[][]{
            {"52", "09", "6A", "D5", "30", "36", "A5", "38", "BF", "40", "A3", "9E", "81", "F3", "D7", "FB"},
            {"7C", "E3", "39", "82", "9B", "2F", "FF", "87", "34", "8E", "43", "44", "C4", "DE", "E9", "CB"},
            {"54", "7B", "94", "32", "A6", "C2", "23", "3D", "EE", "4C", "95", "0B", "42", "FA", "C3", "4E"},
            {"08", "2E", "A1", "66", "28", "D9", "24", "B2", "76", "5B", "A2", "49", "6D", "8B", "D1", "25"},
            {"72", "F8", "F6", "64", "86", "68", "98", "16", "D4", "A4", "5C", "CC", "5D", "65", "B6", "92"},
            {"6C", "70", "48", "50", "FD", "ED", "B9", "DA", "5E", "15", "46", "57", "A7", "8D", "9D", "84"},
            {"90", "D8", "AB", "00", "8C", "BC", "D3", "0A", "F7", "E4", "58", "05", "B8", "B3", "45", "06"},
            {"D0", "2C", "1E", "8F", "CA", "3F", "0F", "02", "C1", "AF", "BD", "03", "01", "13", "8A", "6B"},
            {"3A", "91", "11", "41", "4F", "67", "DC", "EA", "97", "F2", "CF", "CE", "F0", "B4", "E6", "73"},
            {"96", "AC", "74", "22", "E7", "AD", "35", "85", "E2", "F9", "37", "E8", "1C", "75", "DF", "6E"},
            {"47", "F1", "1A", "71", "1D", "29", "C5", "89", "6F", "B7", "62", "0E", "AA", "18", "BE", "1B"},
            {"FC", "56", "3E", "4B", "C6", "D2", "79", "20", "9A", "DB", "C0", "FE", "78", "CD", "5A", "F4"},
            {"1F", "DD", "A8", "33", "88", "07", "C7", "31", "B1", "12", "10", "59", "27", "80", "EC", "5F"},
            {"60", "51", "7F", "A9", "19", "B5", "4A", "0D", "2D", "E5", "7A", "9F", "93", "C9", "9C", "EF"},
            {"A0", "E0", "3B", "4D", "AE", "2A", "F5", "B0", "C8", "EB", "BB", "3C", "83", "53", "99", "61"},
            {"17", "2B", "04", "7E", "BA", "77", "D6", "26", "E1", "69", "14", "63", "55", "21", "0C", "7D"}
    };

    private static final byte[][] MixColumnMatrix = new byte[][]{
            {0x02, 0x03, 0x01, 0x01},
            {0x01, 0x02, 0x03, 0x01},
            {0x01, 0x01, 0x02, 0x03},
            {0x03, 0x01, 0x01, 0x02}
    };

    private static final byte[][] InvMixColumnMatrix = new byte[][]{
            {0x0E, 0x0B, 0x0D, 0x09},
            {0x09, 0x0E, 0x0B, 0x0D},
            {0x0D, 0x09, 0x0E, 0x0B},
            {0x0B, 0x0D, 0x09, 0x0E}
    };

    private static final String[] RCONsStr = new String[]{
            "01000000", "02000000", "04000000", "08000000", "10000000",
            "20000000", "40000000", "80000000", "1b000000", "36000000"
    };

    private static final ByteStream[] RCONs = new ByteStream[10];

    private static final ByteGF GFOp = new ByteGF();

    static {
        try {
            for (int i = 0; i < 16; i++) {
                for (int j = 0; j < 16; j++) {
                    SBOXs[i][j] = Hex.decodeHex(SBOXsStr[i][j])[0];
                    InverseSBOXs[i][j] = Hex.decodeHex(InverseSBOXsStr[i][j])[0];
                }
            }
            for (int i = 0; i < 10; i++) {
                RCONs[i] = new ByteStream(Hex.decodeHex(RCONsStr[i]), 4);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 用于将种子密钥生成轮密钥
     *
     * @param initKey 算法的输入为16字节的密钥和加密轮数
     * @param N       轮数
     * @return 输出为轮密钥
     */
    public ByteStream[] generateKey(ByteStream initKey, int N) {

        byte[][] keyBytes = new byte[4 * (N + 1)][4];
        ByteStream[] roundKey = new ByteStream[N + 1];
        try {
            for (int i = 0; i < 4; i++) {
                for (int j = 0; j < 4; j++) {
                    keyBytes[i][j] = initKey.get(i * 4 + j + 1);
                }
            }
            // keyBytes[i][j]，如果i不是4的倍数，那么W[i][]=W[i-4][]^W[i-1][]，
            // 如果i是4的倍数，那么W[i][]=W[i-4][]^T(W[i-1][])
            byte[] temp = null;
            for (int i = 4; i < 4 * (N + 1); i++) {
                if (i % 4 != 0) {
                    temp = keyBytes[i - 1];
                } else {
                    temp = T(keyBytes[i - 1], i / 4);
                }
                for (int j = 0; j < 4; j++) {
                    keyBytes[i][j] = (byte) (keyBytes[i - 4][j] ^ temp[j]);
                }
            }
            // 用这4*(N+1)个字生成N+1个轮密钥
            for (int round = 0; round < N + 1; round++) {
                roundKey[round] = new ByteStream(16);
                for (int i = 0; i < 4; i++) {
                    for (int j = 0; j < 4; j++) {
                        roundKey[round].set(i * 4 + j + 1, keyBytes[round * 4 + i][j]);
                    }
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return roundKey;
    }


    /**
     * 封装一个T函数
     *
     * @param word 输入的1个字，也即4字节
     * @param j    轮数
     * @return 返回经过处理的1个字，一共4个字节
     */
    private byte[] T(byte[] word, int j) throws Exception {
        // 1. 字循环，输入的4个字节，循环向左移动1个字节，即[0,1,2,3]变成[1,2,3,0]
        byte[] res = new byte[4];
        res[0] = word[1];
        res[1] = word[2];
        res[2] = word[3];
        res[3] = word[0];
        // 2. 字代换，使用SBOXs进行代换
        for (int i = 0; i < 4; i++) {
            int indexByte = Byte.toUnsignedInt(res[i]);
            res[i] = SBOXs[indexByte / 16][indexByte & 0x0f];
        }
        // 3. 轮常量异或
        for (int i = 0; i < 3; i++) {
            res[i] = (byte) (res[i] ^ RCONs[j - 1].get(i + 1));
        }
        return res;
    }

    public ByteStream encrypt(ByteStream plaintext, ByteStream[] roundKey, int N) {
        // 初始轮
        ByteStream state = addKey(plaintext, roundKey[0]);
        // ByteStream addKey = MiddleData.encryptMiddleDatas[0].getAddKey();
        // System.out.println(state.toHexString() + " / " + addKey.toHexString() + " / " + state.equals(addKey));
        // 进行[1,N-1]轮
        for (int i = 1; i <= N - 1; i++) {
            ByteStream afterByteSub = byteSub(state, false);
            // ByteStream x = MiddleData.encryptMiddleDatas[i].getByteSub();
            // System.out.println(afterByteSub.toHexString() + " / " + x.toHexString() + " / " + afterByteSub.equals(x));

            ByteStream afterShiftRow = shiftRow(afterByteSub, false);
            // x = MiddleData.encryptMiddleDatas[i].getShiftRow();
            // System.out.println(afterShiftRow.toHexString() + " / " + x.toHexString() + " / " + afterShiftRow.equals(x));

            ByteStream afterMixColumn = mixColumn(afterShiftRow, false);
            // x = MiddleData.encryptMiddleDatas[i].getMixColumn();
            // System.out.println(afterMixColumn.toHexString() + " / " + x.toHexString() + " / " + afterMixColumn.equals(x));

            state = addKey(afterMixColumn, roundKey[i]);
            // x = MiddleData.encryptMiddleDatas[i].getAddKey();
            // System.out.println(state.toHexString() + " / " + x.toHexString() + " / " + state.equals(x));
        }
        // 进行第N轮，也就是最后1轮
        ByteStream afterByteSub = byteSub(state, false);
        // ByteStream x = MiddleData.encryptMiddleDatas[N].getByteSub();
        // System.out.println(afterByteSub.toHexString() + " / " + x.toHexString() + " / " + afterByteSub.equals(x));

        ByteStream afterShiftRow = shiftRow(afterByteSub, false);
        // x = MiddleData.encryptMiddleDatas[N].getShiftRow();
        // System.out.println(afterShiftRow.toHexString() + " / " + x.toHexString() + " / " + afterShiftRow.equals(x));

        ByteStream ciphertext = addKey(afterShiftRow, roundKey[N]);
        // x = MiddleData.encryptMiddleDatas[N].getAddKey();
        // System.out.println(ciphertext.toHexString() + " / " + x.toHexString() + " / " + ciphertext.equals(x));

        return ciphertext;
    }

    public ByteStream decrypt(ByteStream ciphertext, ByteStream[] roundKey, int N) {
        ByteStream state = addKey(ciphertext, roundKey[N]);
        // System.out.println(state.toHexString());
        // 进行N-1轮运算
        for (int i = N - 1; i >= 1; i--) {
            ByteStream afterShiftRow = shiftRow(state, true);
            // System.out.println(afterShiftRow.toHexString());

            ByteStream afterByteSub = byteSub(afterShiftRow, true);
            // System.out.println(afterByteSub.toHexString());

            ByteStream afterAddKey = addKey(afterByteSub, roundKey[i]);
            // System.out.println(afterAddKey.toHexString());

            state = mixColumn(afterAddKey, true);
            // System.out.println(state.toHexString());
        }
        // 最后一轮
        ByteStream afterShiftRow = shiftRow(state, true);
        // System.out.println(afterShiftRow.toHexString());

        ByteStream afterByteSub = byteSub(afterShiftRow, true);
        // System.out.println(afterByteSub.toHexString());

        ByteStream plaintext = addKey(afterByteSub, roundKey[0]);
        // System.out.println(plaintext.toHexString());

        // 返回明文
        return plaintext;
    }

    /**
     * 基本变换函数：轮密钥加
     *
     * @param afterMIxColumn 经过列混淆的消息，总共16字节
     * @param oneKey         一个轮密钥
     * @return 输出处理后的消息
     */
    private ByteStream addKey(ByteStream afterMIxColumn, ByteStream oneKey) {
        ByteStream result = null;
        try {
            result = new ByteStream(16);
            for (int i = 1; i <= 16; i++) {
                result.set(i, (byte) (afterMIxColumn.get(i) ^ oneKey.get(i)));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return result;
    }

    /**
     * 基本变换函数：字节代替
     *
     * @param state   经过轮密钥加处理的消息，总共16字节
     * @param invFlag 用以表示是否是解密过程中的逆运算，加密时其为true，解密时其为false
     * @return 16字节的经过S盒的消息
     */
    private ByteStream byteSub(ByteStream state, boolean invFlag) {
        ByteStream afterByteSub = null;
        byte[][] HERE_BOX = invFlag ? InverseSBOXs : SBOXs;
        try {
            afterByteSub = new ByteStream(16);
            for (int i = 1; i <= 16; i++) {
                int indexByte = Byte.toUnsignedInt(state.get(i));
                afterByteSub.set(i, HERE_BOX[indexByte / 16][indexByte & 0x0f]);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        return afterByteSub;
    }

    /**
     * 基本变换函数：行移位算法
     *
     * @param afterByteSub 输入是经过S盒的处理的消息，一共16字节
     * @param invFlag      用以表示是否是解密过程中的逆运算，加密时其为true，解密时其为false
     * @return 输出经过行移位的消息，一共是16字节
     */
    private ByteStream shiftRow(ByteStream afterByteSub, boolean invFlag) {
        // 16个字节其实是4x4，4行4列
        // 加密时：第i行的向左循环移位i个字节，解密时：第i行的向左循环移位3-i个字节
        int[] shiftBit = invFlag ? new int[]{0, 3, 2, 1} : new int[]{0, 1, 2, 3};
        byte temp = 0;
        for (int i = 1; i <= 4; i++) {
            try {
                temp = afterByteSub.get(i);
                for (int time = 0; time < shiftBit[i - 1]; time++) {
                    afterByteSub.set(i, afterByteSub.get(i + 4));
                    afterByteSub.set(i + 4, afterByteSub.get(i + 8));
                    afterByteSub.set(i + 8, afterByteSub.get(i + 12));
                    afterByteSub.set(i + 12, temp);
                    temp = afterByteSub.get(i);
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        return afterByteSub;
    }

    /**
     * 基本变换函数：列混淆
     *
     * @param afterShiftRow 经过行移位的消息，一共16字节
     * @param invFlag       用以表示是否是解密过程中的逆运算，加密时其为true，解密时其为false
     * @return 返回经过列混淆的消息，一共16字节
     */
    private ByteStream mixColumn(ByteStream afterShiftRow, boolean invFlag) {
        // MixColumnMatrix是列混淆矩阵，InvMixColumnMatrix是逆向列混淆矩阵
        // MixColumnMatrix或InvMixColumnMatrix，与afterShiftRow进行有限域GF(2^8)的矩阵乘法，
        // 得到afterMixColumn或者invAfterMixColumn
        ByteStream afterMixColumn = null;
        byte[][] HERE_MATRIX = invFlag ? InvMixColumnMatrix : MixColumnMatrix;

        // 有限域上的矩阵乘法
        try {
            afterMixColumn = new ByteStream(16);
            for (int i = 1; i <= 4; i++) {
                for (int j = 1; j <= 4; j++) {
                    byte res = GFOp.add(
                            GFOp.mul(HERE_MATRIX[i - 1][0], afterShiftRow.get(1, j)),
                            GFOp.mul(HERE_MATRIX[i - 1][1], afterShiftRow.get(2, j)),
                            GFOp.mul(HERE_MATRIX[i - 1][2], afterShiftRow.get(3, j)),
                            GFOp.mul(HERE_MATRIX[i - 1][3], afterShiftRow.get(4, j))
                    );
                    afterMixColumn.set(i, j, res);
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return afterMixColumn;
    }
}
